import sys
import os
import random
from pox.core import core
from pox.lib.util import dpid_to_str
from pox.lib.addresses import IPAddr, EthAddr
import pox.openflow.libopenflow_01 as of
import pox.lib.packet.arp as arp

"""first-test version(v0.2), it does perform arp packages modification
works with connection to multiple OVS, requires 2 ovs-switches on one HW(one connected, 
second - modifies arp e.g. filter)"""

"""version works with static network, also it has no idle timeout - so what it's learned for the
first time, it will do forever - it's not flexible ofcourse, will be fixed in future versions"""

#TODO
#add idle timeout function

#TODO
#commit rule only when both partners know what are the ports of destination

#TODO
#implement arp unicast learning

#TODO
#fix external port learning

def mac_generator():
	letter = lambda : random.choice('0123456789abcdef')
	octet = lambda : letter()+letter()
	mac = [octet() for i in range(6)]
	mac[0] = 'ff'	#simple requirment for broadcast and localbits
	return EthAddr(':'.join(mac))

log = core.getLogger()
broadcast_mac = EthAddr('ff:ff:ff:ff:ff:ff')
network_mac = mac_generator() #generates random mac_id for current session
log.info('Generated unique mac address: %s' % str(network_mac))

class Host(object):
	def __init__(self, mac=None, ip=None, port=None, domain=None):
		#number of ovs port for the specific host
		self.port = port
		self.ip = ip
		self.mac = mac
		self.domain = domain

class HwNodeOVS(object):
	"""One class to rule all the switches"""
	def __init__(self, connection, topo_data, dpid):
		self.connection = connection
		connection.addListeners(self)	#listens packages from this connection only
		self.EXTERNAL_PORT = None	#port directed to the cloud

		self.ip_list = []
		for cluster in topo_data:	#list of strings
			self.ip_list += cluster

		self.cluster_map = {}	#dictionary {ip:[connected to ip]}, just a helper object
		for cluster in topo_data:
			for host_ip in cluster:
				self.cluster_map[host_ip] = cluster

		self.host_dict = {}	#mapping ip:Host<hw_mac, ip, port, domain>

	def _handle_PacketIn(self,event):
		packet = event.parsed	#packet actually is a base_packet class
		
		if packet.type == packet.ARP_TYPE:
			ip_src = str(packet.payload.protosrc)
			ip_dst = str(packet.payload.protodst)
			mac_src = str(packet.src)
			mac_dst = str(packet.dst)
			if packet.payload.opcode == arp.REQUEST:
				self.__handle_arp_request_packet(event)	#no learning from arp-request
			elif packet.payload.opcode == arp.REPLY:
				if ip_src in self.ip_list and ip_dst in self.cluster_map.get(ip_src):

					for key in [(mac_src,ip_src),(mac_dst, ip_dst)]:
						if key not in self.host_dict.keys():
							self.host_dict[key] = Host(mac=key[0],ip=key[1],domain=self.cluster_map.get(key[1]))
					if (self.host_dict[(mac_src, ip_src)].port == None):
						self.host_dict[(mac_src, ip_src)].port == event.port
					self.__handle_arp_reply_packet(event) #learning from arp-reply

		elif packet.type == packet.IP_TYPE:
			ip_src = str(packet.payload.srcip)
			ip_dst = str(packet.payload.dstip)
			mac_src = str(packet.src)
			mac_dst = str(packet.dst)
			
			if (mac_src,ip_src) in self.host_dict.keys() and (mac_dst,ip_dst) in self.host_dict.keys():	#preventing KeyError exception
				if (ip_dst in self.host_dict[(mac_src,ip_src)].domain and 
				   mac_src == self.host_dict[(mac_src,ip_src)].mac and 
				   mac_dst == self.host_dict[(mac_dst,ip_dst)].mac):					
				   
					for key in [(mac_src,ip_src),(mac_dst, ip_dst)]:
						if key not in self.host_dict.keys():
							self.host_dict[key] = Host(mac=key[0],ip=key[1],domain=self.cluster_map.get(key[1]))
					if (self.host_dict[(mac_src, ip_src)].port == None):
						self.host_dict[(mac_src, ip_src)].port = event.port
					log.info('packetin')
					self.__handle_IPv4_packet(event)
			#ipv4 packets handled

	def __handle_arp_request_packet(self, event):
		if event.parsed.dst == network_mac:	#packet came from the cloud
			if self.EXTERNAL_PORT == None:	#very important-finding the external port
				self.EXTERNAL_PORT = event.port

			msg = of.ofp_flow_mod()	#write corresponding rules
			msg.match.dl_type = 0x0806
			msg.match.dl_dst = network_mac
			msg.actions.append(of.ofp_action_dl_addr.set_dst(broadcast_mac))
			msg.actions.append(of.ofp_action_output(port = of.OFPP_ALL))
			event.connection.send(msg)

			msg = of.ofp_packet_out(data=event.ofp)	#forward incoming packet
			msg.actions.append(of.ofp_action_dl_addr.set_dst(broadcast_mac))	#change network_mac_id to broadcast mac
			msg.actions.append(of.ofp_action_output(port = of.OFPP_ALL))
			event.connection.send(msg)

		elif event.parsed.dst == broadcast_mac: #packet came from the inside host
			if self.EXTERNAL_PORT:
				msg = of.ofp_flow_mod()	#write corresponding rules
				msg.match.dl_type = 0x0806
				msg.match.dl_dst = broadcast_mac
				msg.actions.append(of.ofp_action_output(port = of.OFPP_ALL))
				msg.actions.append(of.ofp_action_dl_addr.set_dst(network_mac))
				msg.actions.append(of.ofp_action_output(port = self.EXTERNAL_PORT))
				event.connection.send(msg)
			else:
				msg = of.ofp_packet_out(data=event.ofp)	#while don't know the external port - forward to all ports
				msg.actions.append(of.ofp_action_output(port = of.OFPP_ALL))
				msg.actions.append(of.ofp_action_dl_addr.set_dst(network_mac))		
				msg.actions.append(of.ofp_action_output(port = of.OFPP_ALL))
				event.connection.send(msg)

	def __handle_arp_reply_packet(self, event):
		mac_dst = str(event.parsed.dst)
		mac_src = str(event.parsed.dst)
		ip_dst = str(event.parsed.payload.protodst)
		ip_src = str(event.parsed.payload.protosrc)
		host_dst = self.host_dict.get((mac_dst,ip_dst)) 
		if host_dst and host_dst.port: #required port found
			msg = of.ofp_flow_mod()	#add-flow instance to the flow table
			msg.match.dl_type = 0x0806
			msg.match.dl_src = event.parsed.src
			msg.match.dl_dst = event.parsed.dst
			msg.match.nw_src = event.parsed.payload.protosrc
			msg.match.nw_dst = event.parsed.payload.protodst
			msg.actions.append(of.ofp_action_output(port=host_dst.port))
			event.connection.send(msg)

			msg = of.ofp_flow_mod()	#write reversed arp unicast flow to switch
			msg.match.dl_type = 0x0806
			msg.match.dl_src = event.parsed.dst
			msg.match.dl_dst = event.parsed.src
			msg.match.nw_src = event.parsed.payload.protodst
			msg.match.nw_dst = event.parsed.payload.protosrc
			msg.actions.append(of.ofp_action_output(port=event.port))
			event.connection.send(msg)

			msg = of.ofp_packet_out(data=event.ofp)	#forward the incoming packet
			msg.actions.append(of.ofp_action_output(port=host_dst.port))
			event.connection.send(msg)
		else:
			msg = of.ofp_packet_out(data=event.ofp)	#forward the incoming packet(no rule writing)
			msg.actions.append(of.ofp_action_output(port=of.OFPP_ALL)) #sending arp-reply message to all ports //bad decision
			event.connection.send(msg)

	def __handle_IPv4_packet(self, event):
		mac_dst = str(event.parsed.dst)
		mac_src = str(event.parsed.src)
		ip_dst = str(event.parsed.payload.dstip)
		ip_src = str(event.parsed.payload.srcip)
		host_dst = self.host_dict.get((mac_dst,ip_dst))
		if host_dst and host_dst.port:
			msg = of.ofp_flow_mod()	#add flow-item to flow table
			msg.match.dl_type = 0x0800	#tcp type package
			msg.match.dl_src = event.parsed.src
			msg.match.dl_dst = event.parsed.dst
			msg.match.nw_src = event.parsed.payload.srcip
			msg.match.nw_dst = event.parsed.payload.dstip
			msg.actions.append(of.ofp_action_output(port=host_dst.port))
			event.connection.send(msg)

			msg = of.ofp_flow_mod()	#reversed flow
			msg.match.dl_type = 0x0800
			msg.match.dl_src = event.parsed.dst
			msg.match.dl_dst = event.parsed.src
			msg.match.nw_src = event.parsed.payload.dstip
			msg.match.nw_dst = event.parsed.payload.srcip
			msg.actions.append(of.ofp_action_output(port=event.port))
			event.connection.send(msg)

			msg = of.ofp_packet_out(data=event.ofp)	#forward the incoming packet
			msg.actions.append(of.ofp_action_output(port=host_dst.port))
			event.connection.send(msg)
		else:
			msg = of.ofp_packet_out(data=event.ofp)	#forward the incoming packet
			msg.actions.append(of.ofp_action_output(port=of.OFPP_ALL))
			event.connection.send(msg)

class Setup(object):
	"""Listens for incoming connections from ovs and sends all logic
	and connection data to HwNodeSwitch class"""
	def __init__(self, inputfile):
		#listen to openflow connection
		core.openflow.addListeners(self)
		if not (self.__parse_input(inputfile)):
			sys.exit()
		log.info('Inputfile %s -- private networks data parsed.' % (inputfile))
		self.connection_list = []

	def _handle_ConnectionUp(self, event):
		log.info('Catched connection:%s' % (event.connection))
		self.connection_list.append(HwNodeOVS(event.connection, self.topo_data, event.dpid))	
		#added to list, to be sure that garbage collector won't clean it down

	def __parse_input(self, inputfile):
		try:
			fsock = open(inputfile, 'r')
			content = fsock.readlines()
			self.topo_data = [line.rstrip('\n').split(',') for line in content]
			fsock.close()
		except IOError:
			log.fatal('Can\'t get access to input file: %s' % (inputfile))
			return False
		return True

def print_usage():
	print """Usage: %s --inputfile=file""" % ('containers_cloud')

def launch(inputfile=None):
	if inputfile:
		core.registerNew(Setup, inputfile)
	else:
		log.fatal('No file input. Exiting...')